# 推荐模型
import torch
from torch import nn
from torch.utils.data import DataLoader
from torchvision import datasets, transforms

class RecModel(nn.Module):
    def __init__(self):
        self.name = 'rec_model.RecModel'

    @staticmethod
    def demo():
        # 题目知识点向量
        # q_kpv = torch.zeros((2,2), dtype=torch.float32)
        # u_kpv = torch.zeros((5,2), dtype=torch.float32)
        # print(f'q_kpv: {q_kpv.shape}; ?????????????')
        # s1 = torch.dot(q_kpv, u_kpv)
        # print(f's1: {s1}; {type(s1)}; {type(s1.detach().cpu().item())}; {s1.detach().cpu().item()};')

        D_kp = 2
        Nq = 5 # 题目数量
        Nu = 4 # 学生数量
        Y = torch.tensor([
            [5.0, 5.0, 0.0, 0.0],
            [5.0, 0.0, 0.0, 0.0],
            [0.0, 4.0, 0.0, 0.0],
            [0.0, 0.0, 5.0, 4.0],
            [0.0, 0.0, 5.0, 0.0]
        ], dtype=torch.float32)
        mask_Y = torch.tensor([
            [1.0, 1.0, 1.0, 1.0],
            [1.0, 0.0, 0.0, 1.0],
            [0.0, 1.0, 1.0, 0.0],
            [1.0, 1.0, 1.0, 1.0],
            [1.0, 1.0, 1.0, 0.0]
        ], dtype=torch.float32)
        Y_ = Y * mask_Y
        # X = torch.randn((Nq, D_kp), dtype=torch.float32, requires_grad=True)
        X = torch.tensor([
            [0.9, 0.0],
            [1.0, 0.01],
            [0.99, 0.0],
            [0.1, 1.0],
            [0.0, 0.9]
        ], dtype=torch.float32, requires_grad=True)
        # Theta = torch.randn((Nu, D_kp), dtype=torch.float32, requires_grad=True)
        Theta = torch.tensor([
            [5.0, 0.0],
            [5.0, 0.01],
            [0.0, 5.0],
            [0.0, 5.0]
        ], dtype=torch.float32, requires_grad=True)
        model_params = [X, Theta]
        optimizer = torch.optim.SGD(model_params, lr=1e-5)
        epochs = 5000000
        for epoch in range(epochs):
            optimizer.zero_grad()
            Y_hat_raw = X @ Theta.T
            Y_hat = Y_hat_raw * mask_Y
            print(f'Y_hat:\n{Y_hat};')
            loss = torch.norm(Y_hat - Y)
            if epoch % 1000 == 0:
                print(f'epoch_{epoch}: loss: {type(loss)}; {loss};')
            loss.backward()
            optimizer.step()
            # 限制参数的范围
            X.data.clamp_(min=0.0, max=1.0)
            Theta.data.clamp_(min=0.0, max=5.0)
            # Early stopping
            if loss.detach().cpu().item() < 0.001:
                break
        print(f'X:\n{X};')
        print(f'Theta:\n{Theta};')
